package main

import (
	"context"
	"crypto/md5"
	"log"
	"navidb-master/config"
	"navidb-master/global"
	"navidb-master/logger"
	"sync"
	"time"

	pb "navidb-master/proto"

	"github.com/panjf2000/ants/v2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
)

var (
	keepaliveParams = keepalive.ClientParameters{
		Time:                10 * time.Second, // 发送心跳的间隔
		Timeout:             1 * time.Second,  // 等待ack的超时时间
		PermitWithoutStream: true,             // 没有活跃流时也发送心跳
	}
	clients = make(map[int]*pb.NavidbServiceClient)
	servers = make(map[int]string)
)

type TaskData struct {
	index int
	key   string
	wg    *sync.WaitGroup
}

func main() {
	config.Init()
	logger.Init()

	for index, item := range config.Conf.Nodes.Address {
		logger.Infof("<<< connecting to node [%s]", item)
		conn, err := grpc.NewClient(
			item,
			grpc.WithTransportCredentials(insecure.NewCredentials()),
		)
		if err != nil {
			logger.Fatalf("did not connect: %v", err)
		}
		defer conn.Close()
		client := pb.NewNavidbServiceClient(conn)
		clients[index] = &client
		servers[index] = item
	}
	global.Nodes = clients

	var sw sync.WaitGroup

	key := "2e65029dcb7a861b3f7d1098da6004be"
	start := time.Now().Unix()

	pool, _ := ants.NewPoolWithFunc(1000, func(a any) {
		data := a.(*TaskData)
		getResp, err := (*GetShard(key)).Get(context.TODO(), &pb.GetDataReq{Key: key})
		if err != nil {
			log.Fatalf("could not get: %v", err)
		}
		log.Printf("Get %d response: %s", data.index, getResp.GetValue())
		data.wg.Done()
	})
	defer pool.Release()

	for i := range 2000 {
		sw.Add(1)
		pool.Invoke(&TaskData{wg: &sw, index: i})
	}
	sw.Wait()
	log.Printf("cost time: %d", time.Now().Unix()-start)
}

func GetShard(key string) *pb.NavidbServiceClient {
	hash := md5.Sum([]byte(key))
	shardId := int(hash[0]) % len(config.Conf.Nodes.Address)
	return global.Nodes[shardId]
}
